import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
import random


class Brown(object):
    """
    布朗运动
    S(t) = S(0) * exp((mu - signa^2 / 2) * t + sigma * time)
    """
    def __init__(self):
        # matplotlib中文显示方块
        mpl.rcParams['font.sans-serif'] = ['SimHei']  # 指定默认字体
        mpl.rcParams['axes.unicode_minus'] = False  # 解决保存图像是负号'-'显示为方块的问题

    @staticmethod
    def brown_base():
        """
        布朗运动（随机游走）
        :return:
        """
        rect = [0.1, 5.0, 0.1, 0.1]
        fig = plt.figure(figsize=(10, 10))

        T = 2
        mu = 0.1
        sigma = 0.04
        s0 = 20
        dt = 0.01
        N = round(T / dt)
        t = np.linspace(0, T, N)

        w = np.random.standard_normal(size=N)
        w = np.cumsum(w) * np.sqrt(dt)

        # 指数，类似正态分布
        x = (mu - 0.5 * np.power(sigma, 2)) * t + sigma * w

        # 布朗运动（正态分布公式）
        s = s0 * np.exp(x)

        plt.plot(t, s)
        plt.xlabel('时间 t')
        plt.ylabel('位置 s')
        plt.title('布朗运动（随机游走）')
        plt.show()


class BrownStock(object):
    """
    布朗运动模拟股票市场价格波动
    """
    def __init__(self, price, distribution):
        self.price = price
        self.distribute = distribution
        self.history = [price]
        self.last_change = 0
        # matplotlib中文显示方块
        mpl.rcParams['font.sans-serif'] = ['SimHei']  # 指定默认字体
        mpl.rcParams['axes.unicode_minus'] = False  # 解决保存图像是负号'-'显示为方块的问题

    def get_price(self):
        return self.price

    def set_price(self, price):
        self.price = price
        self.history.append(price)

    def walk_it(self, marketbias, mo):
        """
        随机游走
        :param marketbias:
        :param mo:
        :return:
        """
        old_price = self.price
        base_move = self.distribute() + marketbias
        self.price = self.price * (1 + base_move)
        if mo:
            # random.gauss(0.5, 0.5) 生成符合条件的随机数
            self.price += random.gauss(0.5, 0.5) * self.last_change

        self.price = 0 if self.price < 0.01 else self.price
        self.history.append(self.price)
        self.last_change = self.price - old_price

    def plot_it(self):
        plt.plot(self.history)
        plt.title('布朗运动模拟股票市场价格波动')
        plt.xlabel('日期')
        plt.ylabel('价格')


def brown_stock_main():
    def run_simulation(stocks, bias, mo, num_days):
        mean = 0.0
        for s in stocks:
            for d in range(num_days):
                s.walk_it(bias, mo)
            s.plot_it()
            mean += s.get_price()
        mean /= float(num_stocks)
        plt.axhline(mean)

    num_stocks = 3
    num_days = 400
    stocks = []
    bias = 0
    mo = False
    startvalues = [107, 89, 37.8]
    for value in startvalues:
        # 均匀分布采样：涨跌幅 <= 10%
        volatility = random.uniform(0, 0.1)
        # 均匀分布
        d1 = lambda : random.uniform(-volatility, volatility)
        stocks.append(BrownStock(value, d1))

    run_simulation(stocks, bias, mo, num_days)

    plt.show()


# brown = Brown()
# 布朗运动
# brown.brown_base()

brown_stock_main()
